package com.yuma.springaiollama.controller;

import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.content.Media;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

@RestController
public class OllamaController {
    @Autowired
    private OllamaChatModel ollamaChatModel;
    @GetMapping("/ollama/call/{msg}")
    public String call(@PathVariable String msg) {
        String call = ollamaChatModel.call(msg);
        System.out.println(call);
        return call;
    }
    // 流式输出
    @GetMapping(value = "/ollama/stream/{msg}",produces = "text/html;charset=utf-8")
    public Flux<String> stream(@PathVariable String msg) {
        return ollamaChatModel.stream(msg);
    }

    @GetMapping("/ollama/multiModel/{msg}")
    public String multiModel(@PathVariable String msg) {
        var imageResource = new ClassPathResource("/img.jpeg");
        OllamaOptions ollamaOptions = OllamaOptions.builder()
                .model("gemma3:4b")
                .build();
        Media media = new Media(MimeTypeUtils.IMAGE_JPEG, imageResource);
        ChatResponse response = ollamaChatModel.call(
                new Prompt(
                        UserMessage.builder().media(media).text(msg).build(),
                        ollamaOptions
                )
        );
        return response.getResult().getOutput().getText();
    }
}
